from pyspark.sql import SparkSession
from pyspark.sql.functions import col, avg, max, min, count, round, lag, datediff
from pyspark.sql.window import Window
import matplotlib.pyplot as plt
import pandas as pd
import os
from matplotlib import font_manager
import seaborn as sns
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import LinearRegression
from datetime import datetime, timedelta

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS的通用中文字体
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置数据目录的绝对路径
DATA_DIR = "/Users/guopeiran/Code/VSCode/Something/something/stockdata"

# 创建SparkSession
spark = SparkSession.builder \
    .appName("Stock Analysis") \
    .getOrCreate()

# 读取所有股票数据
stock_data = spark.read.csv(os.path.join(DATA_DIR, "*.csv"), header=True, inferSchema=True)

def plot_simple_price_trend(stock_code):
    """绘制简化的价格走势图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    plt.plot(range(len(stock_pdf)), stock_pdf['close'], label='收盘价', linewidth=2, color='#2ecc71')
    
    # 添加趋势说明
    first_price = stock_pdf['close'].iloc[0]
    last_price = stock_pdf['close'].iloc[-1]
    change_percent = ((last_price - first_price) / first_price) * 100
    
    trend_text = f"整体趋势: {'上涨' if change_percent > 0 else '下跌'} {abs(change_percent):.1f}%"
    plt.text(0.02, 0.98, trend_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票价格走势图', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('价格', fontsize=12)
    
    # 优化日期标签显示
    n_labels = 10
    indices = np.linspace(0, len(stock_pdf)-1, n_labels, dtype=int)
    plt.xticks(indices, stock_pdf['date'].iloc[indices], rotation=45, ha='right')
    
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_price_trend.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_volume(stock_code):
    """绘制简化的成交量图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    # 根据涨跌设置颜色
    colors = ['#e74c3c' if x > 0 else '#2ecc71' for x in stock_pdf['p_change']]
    plt.bar(range(len(stock_pdf)), stock_pdf['volume'], alpha=0.6, color=colors)
    
    # 添加成交量说明
    avg_volume = stock_pdf['volume'].mean()
    volume_text = f"平均成交量: {avg_volume:,.0f}"
    plt.text(0.02, 0.98, volume_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票成交量分析\n(红色表示上涨，绿色表示下跌)', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('成交量', fontsize=12)
    
    # 优化日期标签显示
    n_labels = 10
    indices = np.linspace(0, len(stock_pdf)-1, n_labels, dtype=int)
    plt.xticks(indices, stock_pdf['date'].iloc[indices], rotation=45, ha='right')
    
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_volume_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_price_distribution(stock_code):
    """绘制简化的价格分布图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(10, 6))
    sns.histplot(data=stock_pdf, x='close', bins=30, color='#3498db')
    
    # 添加价格说明
    price_text = f"平均价格: {stock_pdf['close'].mean():.2f}\n最高价: {stock_pdf['close'].max():.2f}\n最低价: {stock_pdf['close'].min():.2f}"
    plt.text(0.02, 0.98, price_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票收盘价分布', fontsize=14, pad=15)
    plt.xlabel('价格', fontsize=12)
    plt.ylabel('出现次数', fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_price_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_trend(stock_code):
    """绘制简化的趋势图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    # 计算5日和20日均线
    stock_pdf['MA5'] = stock_pdf['close'].rolling(window=5).mean()
    stock_pdf['MA20'] = stock_pdf['close'].rolling(window=20).mean()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    plt.plot(range(len(stock_pdf)), stock_pdf['close'], label='收盘价', linewidth=2, color='#2ecc71')
    plt.plot(range(len(stock_pdf)), stock_pdf['MA5'], label='5日均线', linewidth=1, color='#e74c3c')
    plt.plot(range(len(stock_pdf)), stock_pdf['MA20'], label='20日均线', linewidth=1, color='#3498db')
    
    # 添加趋势说明
    trend_text = "趋势说明:\n- 绿线：每日收盘价\n- 红线：5日均线（短期趋势）\n- 蓝线：20日均线（长期趋势）"
    plt.text(0.02, 0.98, trend_text, transform=plt.gca().transAxes, 
             fontsize=10, verticalalignment='top')
    
    plt.title(f'{stock_code}股票趋势分析', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('价格', fontsize=12)
    
    # 优化日期标签显示
    n_labels = 10
    indices = np.linspace(0, len(stock_pdf)-1, n_labels, dtype=int)
    plt.xticks(indices, stock_pdf['date'].iloc[indices], rotation=45, ha='right')
    
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_trend_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

def calculate_technical_indicators(stock_pdf):
    """计算技术指标"""
    # RSI指标 (14天)
    delta = stock_pdf['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    stock_pdf['RSI'] = 100 - (100 / (1 + rs))
    
    # MACD指标
    exp1 = stock_pdf['close'].ewm(span=12, adjust=False).mean()
    exp2 = stock_pdf['close'].ewm(span=26, adjust=False).mean()
    stock_pdf['MACD'] = exp1 - exp2
    stock_pdf['Signal_Line'] = stock_pdf['MACD'].ewm(span=9, adjust=False).mean()
    stock_pdf['MACD_Histogram'] = stock_pdf['MACD'] - stock_pdf['Signal_Line']
    
    # KDJ指标
    low_min = stock_pdf['low'].rolling(window=9).min()
    high_max = stock_pdf['high'].rolling(window=9).max()
    stock_pdf['K'] = 100 * ((stock_pdf['close'] - low_min) / (high_max - low_min))
    stock_pdf['D'] = stock_pdf['K'].rolling(window=3).mean()
    stock_pdf['J'] = 3 * stock_pdf['K'] - 2 * stock_pdf['D']
    
    return stock_pdf

def plot_technical_indicators(stock_code):
    """绘制技术指标图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    stock_pdf = calculate_technical_indicators(stock_pdf)
    
    # 创建子图
    fig = plt.figure(figsize=(15, 12))
    gs = fig.add_gridspec(4, 1, height_ratios=[2, 1, 1, 1])
    
    # 1. 价格和MACD
    ax1 = fig.add_subplot(gs[0])
    ax1.plot(stock_pdf.index, stock_pdf['close'], label='收盘价', color='#2ecc71')
    ax1.set_title(f'{stock_code}技术指标分析', fontsize=14, pad=15)
    ax1.set_ylabel('价格', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.legend(loc='upper left')
    
    # 2. MACD
    ax2 = fig.add_subplot(gs[1])
    ax2.plot(stock_pdf.index, stock_pdf['MACD'], label='MACD', color='#3498db')
    ax2.plot(stock_pdf.index, stock_pdf['Signal_Line'], label='信号线', color='#e74c3c')
    ax2.bar(stock_pdf.index, stock_pdf['MACD_Histogram'], label='MACD柱状', color='#95a5a6', alpha=0.5)
    ax2.set_ylabel('MACD', fontsize=12)
    ax2.grid(True, linestyle='--', alpha=0.7)
    ax2.legend(loc='upper left')
    
    # 3. RSI
    ax3 = fig.add_subplot(gs[2])
    ax3.plot(stock_pdf.index, stock_pdf['RSI'], label='RSI', color='#9b59b6')
    ax3.axhline(y=70, color='#e74c3c', linestyle='--', alpha=0.7)
    ax3.axhline(y=30, color='#2ecc71', linestyle='--', alpha=0.7)
    ax3.set_ylabel('RSI', fontsize=12)
    ax3.grid(True, linestyle='--', alpha=0.7)
    ax3.legend(loc='upper left')
    
    # 4. KDJ
    ax4 = fig.add_subplot(gs[3])
    ax4.plot(stock_pdf.index, stock_pdf['K'], label='K', color='#3498db')
    ax4.plot(stock_pdf.index, stock_pdf['D'], label='D', color='#e74c3c')
    ax4.plot(stock_pdf.index, stock_pdf['J'], label='J', color='#2ecc71')
    ax4.set_ylabel('KDJ', fontsize=12)
    ax4.grid(True, linestyle='--', alpha=0.7)
    ax4.legend(loc='upper left')
    
    # 优化日期标签显示
    for ax in [ax1, ax2, ax3, ax4]:
        n_labels = 10
        indices = np.linspace(0, len(stock_pdf)-1, n_labels, dtype=int)
        ax.set_xticks(indices)
        ax.set_xticklabels(stock_pdf['date'].iloc[indices], rotation=45, ha='right')
    
    plt.tight_layout()
    plt.savefig(f'{stock_code}_technical_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

def predict_stock_price(stock_code, days=30):
    """预测股票价格"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    # 准备数据
    data = stock_pdf[['close']].values
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(data)
    
    # 创建训练数据
    X = []
    y = []
    for i in range(60, len(scaled_data)):
        X.append(scaled_data[i-60:i, 0])
        y.append(scaled_data[i, 0])
    X = np.array(X)
    y = np.array(y)
    
    # 训练模型
    model = LinearRegression()
    model.fit(X, y)
    
    # 预测未来价格
    last_60_days = scaled_data[-60:]
    future_predictions = []
    current_batch = last_60_days.reshape(1, -1)[0]
    
    for _ in range(days):
        prediction = model.predict(current_batch.reshape(1, -1))
        future_predictions.append(prediction[0])
        current_batch = np.append(current_batch[1:], prediction)
    
    # 转换回实际价格
    future_predictions = np.array(future_predictions).reshape(-1, 1)
    future_predictions = scaler.inverse_transform(future_predictions)
    
    # 绘制预测图
    plt.figure(figsize=(15, 6))
    plt.plot(stock_pdf.index[-60:], stock_pdf['close'].iloc[-60:], label='历史价格', color='#2ecc71')
    plt.plot(range(len(stock_pdf)-1, len(stock_pdf)+days-1), future_predictions, label='预测价格', color='#e74c3c', linestyle='--')
    
    plt.title(f'{stock_code}股票价格预测（未来{days}天）', fontsize=14, pad=15)
    plt.xlabel('时间', fontsize=12)
    plt.ylabel('价格', fontsize=12)
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_price_prediction.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    return future_predictions[-1][0]

def filter_stocks(min_price=None, max_price=None, min_volume=None, max_volatility=None):
    """筛选股票"""
    # 首先计算波动率
    window_spec = Window.partitionBy("date")
    stock_volatility = stock_data.withColumn(
        "volatility",
        ((max("high").over(window_spec) - min("low").over(window_spec)) / avg("close").over(window_spec) * 100)
    )
    
    # 应用筛选条件
    query = stock_volatility
    
    if min_price is not None:
        query = query.filter(col("close") >= min_price)
    if max_price is not None:
        query = query.filter(col("close") <= max_price)
    if min_volume is not None:
        query = query.filter(col("volume") >= min_volume)
    if max_volatility is not None:
        query = query.filter(col("volatility") <= max_volatility)
    
    return query

def analyze_single_stock(stock_code):
    """分析单支股票并生成图表"""
    # 生成基础分析图表
    plot_simple_price_trend(stock_code)
    plot_simple_volume(stock_code)
    plot_simple_price_distribution(stock_code)
    plot_simple_trend(stock_code)
    
    # 生成技术指标分析图表
    plot_technical_indicators(stock_code)
    
    # 预测未来价格
    future_price = predict_stock_price(stock_code)
    print(f"\n预测{stock_code}股票30天后的价格为: {future_price:.2f}")

def analyze_all_stocks():
    """分析所有股票并找出最佳投资机会"""
    # 筛选股票
    filtered_stocks = filter_stocks(
        min_price=5.0,         # 最低价格
        max_price=200.0,       # 最高价格
        min_volume=100000,     # 最低成交量
        max_volatility=100.0   # 最大波动率
    )
    
    # 按日期分组计算平均波动率
    stock_volatility = filtered_stocks.groupBy("date").agg(
        avg("volatility").alias("avg_volatility"),
        count("*").alias("stock_count")
    )
    
    # 找出波动率最小的日期
    best_stock = stock_volatility.orderBy("avg_volatility").first()
    
    if best_stock is None:
        print("警告：没有找到符合条件的股票，使用默认分析方法。")
        # 使用默认方法分析所有股票
        stock_volatility = stock_data.groupBy("date").agg(
            ((max("high") - min("low")) / avg("close") * 100).alias("avg_volatility"),
            count("*").alias("stock_count")
        )
        best_stock = stock_volatility.orderBy("avg_volatility").first()
    
    return best_stock

def plot_simple_summary(best_stock_date):
    """生成简化的综合分析报告"""
    # 获取推荐日期的所有股票数据
    recommended_stocks = stock_data.filter(col("date") == best_stock_date)
    recommended_pdf = recommended_stocks.toPandas()
    
    # 创建子图布局
    fig = plt.figure(figsize=(15, 10))
    gs = fig.add_gridspec(2, 2)
    
    # 1. 价格分布图
    ax1 = fig.add_subplot(gs[0, 0])
    sns.histplot(data=recommended_pdf, x='close', bins=30, color='#3498db', ax=ax1)
    ax1.set_title('推荐日期股票价格分布', fontsize=12, pad=10)
    ax1.set_xlabel('价格')
    ax1.set_ylabel('股票数量')
    ax1.grid(True, linestyle='--', alpha=0.7)
    
    # 2. 涨跌分布图
    ax2 = fig.add_subplot(gs[0, 1])
    colors = ['#e74c3c' if x > 0 else '#2ecc71' for x in recommended_pdf['p_change']]
    ax2.bar(range(len(recommended_pdf)), recommended_pdf['p_change'], color=colors)
    ax2.set_title('推荐日期股票涨跌分布\n(红色表示上涨，绿色表示下跌)', fontsize=12, pad=10)
    ax2.set_xlabel('股票序号')
    ax2.set_ylabel('涨跌幅(%)')
    ax2.grid(True, linestyle='--', alpha=0.7)
    
    # 3. 成交量分布图
    ax3 = fig.add_subplot(gs[1, 0])
    sns.histplot(data=recommended_pdf, x='volume', bins=30, color='#9b59b6', ax=ax3)
    ax3.set_title('推荐日期股票成交量分布', fontsize=12, pad=10)
    ax3.set_xlabel('成交量')
    ax3.set_ylabel('股票数量')
    ax3.grid(True, linestyle='--', alpha=0.7)
    
    # 4. 关键指标统计
    ax4 = fig.add_subplot(gs[1, 1])
    ax4.axis('off')
    stats_text = f"""
    推荐日期: {best_stock_date}
    
    市场概况:
    - 股票数量: {len(recommended_pdf)}支
    - 平均价格: {recommended_pdf['close'].mean():.2f}元
    - 最高价格: {recommended_pdf['high'].max():.2f}元
    - 最低价格: {recommended_pdf['low'].min():.2f}元
    
    涨跌情况:
    - 平均涨跌幅: {recommended_pdf['p_change'].mean():.2f}%
    - 上涨股票: {len(recommended_pdf[recommended_pdf['p_change'] > 0])}支
    - 下跌股票: {len(recommended_pdf[recommended_pdf['p_change'] < 0])}支
    
    交易情况:
    - 平均成交量: {recommended_pdf['volume'].mean():,.0f}手
    """
    ax4.text(0.1, 0.5, stats_text, fontsize=12, va='center')
    
    # 设置总标题
    fig.suptitle('推荐股票市场分析报告', fontsize=16, y=0.95)
    
    # 调整布局
    plt.tight_layout()
    plt.savefig('recommended_stocks_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

# 主函数
def main():
    # 分析单支股票（以600415为例）
    analyze_single_stock("600415")
    
    # 分析所有股票
    best_stock = analyze_all_stocks()
    
    # 生成推荐股票的综合分析图表
    plot_simple_summary(best_stock['date'])
    
    print("\n推荐投资股票分析：")
    print(f"根据分析，推荐投资日期为 {best_stock['date']} 的股票，")
    print(f"该日期的股票波动率为 {best_stock['avg_volatility']:.2f}%，")
    print("表明该时期市场相对稳定，投资风险较低。")
    
    print("\n已生成以下分析图表：")
    print("1. 600415_price_trend.png - 价格走势图（带趋势说明）")
    print("2. 600415_volume_analysis.png - 成交量分析图（红涨绿跌）")
    print("3. 600415_price_distribution.png - 价格分布图（带统计说明）")
    print("4. 600415_trend_analysis.png - 趋势分析图（带均线说明）")
    print("5. 600415_technical_analysis.png - 技术指标分析")
    print("6. 600415_price_prediction.png - 价格预测分析")
    print("7. recommended_stocks_analysis.png - 市场分析报告")

if __name__ == "__main__":
    main() 